import ast
import json
import javalang
import esprima

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"
EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"

VALID_NODE_TYPES = [
    NODE_TYPE_DIRECTORY,
    NODE_TYPE_FILE,
    NODE_TYPE_CLASS,
    NODE_TYPE_FUNCTION,
]
VALID_EDGE_TYPES = [
    EDGE_TYPE_CONTAINS,
    EDGE_TYPE_INHERITS,
    EDGE_TYPE_INVOKES,
    EDGE_TYPE_IMPORTS,
]


class JavaAnalyzer:
    def __init__(self, filename, source_code):
        self.filename = filename
        self.source_code = source_code
        self.source_lines = source_code.split("\n")
        self.nodes = []
        self.package_name = ""

    def analyze(self, tree):
        if hasattr(tree, "package") and tree.package:
            self.package_name = tree.package.name

        if hasattr(tree, "types"):
            for type_decl in tree.types:
                self._visit_type(type_decl)

        return self.nodes

    def _get_line_number(self, node):
        """安全地获取节点的行号"""
        try:
            if hasattr(node, "position") and node.position:
                return node.position.line
        except:
            pass
        return 0

    def _get_end_line_number(self, node):
        """获取节点的结束行号"""
        try:
            # 对于Java节点，我们需要遍历其子节点来找到最大的行号
            max_line = self._get_line_number(node)
            if hasattr(node, "body") and node.body:
                for child in node.body:
                    child_line = self._get_line_number(child)
                    if child_line > max_line:
                        max_line = child_line
            return max_line if max_line > 0 else None
        except:
            return None

    def _extract_code_block(self, start_line, end_line):
        """提取指定行范围的代码"""
        if start_line <= 0 or start_line > len(self.source_lines):
            return ""

        if end_line is None or end_line <= 0:
            end_line = start_line

        # 调整为0基索引
        start_idx = start_line - 1
        end_idx = min(end_line, len(self.source_lines))

        return "\n".join(self.source_lines[start_idx:end_idx])

    def _visit_type(self, node):
        if isinstance(node, javalang.tree.ClassDeclaration):
            self._visit_class(node)
        elif isinstance(node, javalang.tree.InterfaceDeclaration):
            self._visit_interface(node)

    def _visit_class(self, node):
        class_name = node.name
        if self.package_name:
            full_class_name = f"{self.package_name}.{class_name}"
        else:
            full_class_name = class_name

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        self.nodes.append(
            {
                "name": full_class_name,
                "parent_type": None,
                "type": NODE_TYPE_CLASS,
                "code": code_block if code_block else f"class {class_name}",
                "start_line": start_line,
                "end_line": end_line,
            }
        )

        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, class_name)

    def _visit_interface(self, node):
        interface_name = node.name
        if self.package_name:
            full_interface_name = f"{self.package_name}.{interface_name}"
        else:
            full_interface_name = interface_name

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        self.nodes.append(
            {
                "name": full_interface_name,
                "parent_type": None,
                "type": NODE_TYPE_CLASS,
                "code": code_block if code_block else f"interface {interface_name}",
                "start_line": start_line,
                "end_line": end_line,
            }
        )

        if hasattr(node, "body"):
            for member in node.body:
                if isinstance(member, javalang.tree.MethodDeclaration):
                    self._visit_method(member, interface_name)

    def _visit_method(self, node, parent_class):
        method_name = node.name
        if self.package_name:
            full_method_name = f"{self.package_name}.{parent_class}.{method_name}"
        else:
            full_method_name = f"{parent_class}.{method_name}"

        start_line = self._get_line_number(node)
        end_line = self._get_end_line_number(node)
        code_block = self._extract_code_block(start_line, end_line)

        self.nodes.append(
            {
                "name": full_method_name,
                "parent_type": NODE_TYPE_CLASS,
                "type": NODE_TYPE_FUNCTION,
                "code": code_block if code_block else f"method {method_name}",
                "start_line": start_line,
                "end_line": end_line,
            }
        )


def analyze_file_java(filepath):
    with open(filepath, "r", encoding="utf-8") as file:
        code = file.read()

    try:
        if filepath.endswith(".java"):
            tree = javalang.parse.parse(code)
            analyzer = JavaAnalyzer(filepath, code)
            return analyzer.analyze(tree)
        elif filepath.endswith(".js"):
            tree = esprima.parseScript(code)
            return []  # JavaScript分析器待实现
        else:
            return []
    except Exception as e:
        print(f"Error parsing {filepath}: {e}")
        return []


if __name__ == "__main__":
    file_path = "/data_ext/ref_code/LocAgent/data/data/pp_data/fmxm/run.py"
    file_path = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test/auto-test-b2b/src/test/java/auto/b2b/check/CheckBusinessAccount.java"
    nodes = analyze_file_java(file_path)
    print(len(nodes))
    for node in nodes:
        print(json.dumps(node, ensure_ascii=False, indent=4))
